package com.matrix.framework.core.component;

import com.matrix.framework.core.config.JwtKeyConfig;
import io.jsonwebtoken.SignatureAlgorithm;
import io.jsonwebtoken.security.Keys;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import javax.crypto.SecretKey;
import java.time.LocalDateTime;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.Base64;

/**
 * JWT密钥管理器
 * 支持多密钥轮换机制、配置化和环境隔离
 *
 * Copyright © 海平面工作室 版权所有
 *
 * @Author: Leo
 * @Create: 2024/12/19 10:00
 * @Since 1.2
 */
@Component
public class JwtKeyManager {

    // 密钥存储：keyId -> KeyInfo
    private final Map<String, KeyInfo> keyStore = new ConcurrentHashMap<>();
    
    // 当前活跃密钥ID
    private volatile String currentKeyId;
    
    // 配置
    private final JwtKeyConfig config;

    public JwtKeyManager(JwtKeyConfig config) {
        this.config = config;
        // 初始化时创建第一个密钥
        initializeFirstKey();
    }

    /**
     * 初始化第一个密钥
     */
    private void initializeFirstKey() {
        // 尝试从配置中读取初始密钥
        String initialKey = getInitialKeyFromConfig();
        
        if (StringUtils.hasText(initialKey)) {
            // 使用配置的密钥
            SecretKey secretKey = createKeyFromConfig(initialKey);
            String keyId = generateKeyId();
            KeyInfo keyInfo = new KeyInfo(secretKey, LocalDateTime.now(), 
                                        LocalDateTime.now().plusDays(config.getRotationDays() + config.getOverlapDays()));
            keyStore.put(keyId, keyInfo);
            currentKeyId = keyId;
            System.out.println("使用配置的初始JWT密钥，密钥ID: " + keyId);
        } else {
            // 生成新的密钥
            rotateKey();
        }
    }

    /**
     * 从配置中获取初始密钥
     */
    private String getInitialKeyFromConfig() {
        // 优先级：环境变量 > 配置文件 > 默认值
        String envKey = System.getenv("MATRIX_JWT_INITIAL_KEY");
        if (StringUtils.hasText(envKey)) {
            System.out.println("使用环境变量中的JWT初始密钥");
            return envKey;
        }
        
        // 从配置文件中读取
        String configKey = config.getInitialKey();
        if (StringUtils.hasText(configKey)) {
            System.out.println("使用配置文件中的JWT初始密钥");
            return configKey;
        }
        
        System.out.println("未配置初始密钥，将自动生成新密钥");
        return null; // 返回null表示需要生成新密钥
    }

    /**
     * 从配置字符串创建密钥
     */
    private SecretKey createKeyFromConfig(String keyString) {
        try {
            // 如果是Base64编码的密钥
            if (keyString.length() >= 32) { // 最小密钥长度
                byte[] keyBytes = Base64.getDecoder().decode(keyString);
                return Keys.hmacShaKeyFor(keyBytes);
            } else {
                // 如果密钥太短，使用它作为种子生成密钥
                return Keys.hmacShaKeyFor(keyString.getBytes());
            }
        } catch (Exception e) {
            // 如果解析失败，生成新密钥
            System.err.println("配置的密钥格式错误，将生成新密钥: " + e.getMessage());
            return Keys.secretKeyFor(SignatureAlgorithm.HS256);
        }
    }

    /**
     * 获取当前活跃密钥
     */
    public SecretKey getCurrentKey() {
        KeyInfo keyInfo = keyStore.get(currentKeyId);
        return keyInfo != null ? keyInfo.getSecretKey() : null;
    }

    /**
     * 根据keyId获取密钥
     */
    public SecretKey getKey(String keyId) {
        KeyInfo keyInfo = keyStore.get(keyId);
        return keyInfo != null ? keyInfo.getSecretKey() : null;
    }

    /**
     * 获取当前活跃密钥ID
     */
    public String getCurrentKeyId() {
        return currentKeyId;
    }

    /**
     * 获取所有有效密钥（用于验证）
     */
    public Map<String, SecretKey> getValidKeys() {
        Map<String, SecretKey> validKeys = new HashMap<>();
        LocalDateTime now = LocalDateTime.now();
        
        keyStore.forEach((keyId, keyInfo) -> {
            if (keyInfo.isValid(now)) {
                validKeys.put(keyId, keyInfo.getSecretKey());
            }
        });
        
        return validKeys;
    }

    /**
     * 轮换密钥
     */
    public synchronized void rotateKey() {
        // 清理过期密钥
        cleanupExpiredKeys();
        
        // 检查是否需要轮换
        if (shouldRotateKey()) {
            String newKeyId = generateKeyId();
            SecretKey newKey = Keys.secretKeyFor(SignatureAlgorithm.HS256);
            
            KeyInfo keyInfo = new KeyInfo(newKey, LocalDateTime.now(), 
                                        LocalDateTime.now().plusDays(config.getRotationDays() + config.getOverlapDays()));
            
            keyStore.put(newKeyId, keyInfo);
            currentKeyId = newKeyId;
            
            System.out.println("JWT密钥已轮换，新密钥ID: " + newKeyId);
        }
    }

    /**
     * 检查是否需要轮换密钥
     */
    private boolean shouldRotateKey() {
        if (currentKeyId == null) {
            return true;
        }
        
        KeyInfo currentKeyInfo = keyStore.get(currentKeyId);
        if (currentKeyInfo == null) {
            return true;
        }
        
        // 如果当前密钥即将过期，则轮换
        LocalDateTime rotationTime = currentKeyInfo.getCreatedAt().plusDays(config.getRotationDays());
        return LocalDateTime.now().isAfter(rotationTime);
    }

    /**
     * 清理过期密钥
     */
    private void cleanupExpiredKeys() {
        LocalDateTime now = LocalDateTime.now();
        List<String> expiredKeys = new ArrayList<>();
        
        keyStore.forEach((keyId, keyInfo) -> {
            if (!keyInfo.isValid(now)) {
                expiredKeys.add(keyId);
            }
        });
        
        expiredKeys.forEach(keyStore::remove);
        
        // 如果密钥数量超过限制，删除最旧的
        if (keyStore.size() > config.getMaxKeys()) {
            keyStore.entrySet().stream()
                    .sorted(Comparator.comparing(entry -> entry.getValue().getCreatedAt()))
                    .limit(keyStore.size() - config.getMaxKeys())
                    .map(Map.Entry::getKey)
                    .forEach(keyStore::remove);
        }
    }

    /**
     * 生成密钥ID
     */
    private String generateKeyId() {
        return "key_" + System.currentTimeMillis();
    }

    /**
     * 密钥信息内部类
     */
    private static class KeyInfo {
        private final SecretKey secretKey;
        private final LocalDateTime createdAt;
        private final LocalDateTime expiresAt;

        public KeyInfo(SecretKey secretKey, LocalDateTime createdAt, LocalDateTime expiresAt) {
            this.secretKey = secretKey;
            this.createdAt = createdAt;
            this.expiresAt = expiresAt;
        }

        public SecretKey getSecretKey() {
            return secretKey;
        }

        public LocalDateTime getCreatedAt() {
            return createdAt;
        }

        public LocalDateTime getExpiresAt() {
            return expiresAt;
        }

        public boolean isValid(LocalDateTime now) {
            return now.isBefore(expiresAt);
        }
    }
} 